import pandas as pd
import numpy as np
from matplotlib import pyplot as plt
from matplotlib.image import imread
import matplotlib
import cv2
from scipy.ndimage import label, find_objects
from scipy.ndimage import label, find_objects, binary_fill_holes
from skimage.measure import find_contours
from matplotlib.patches import Polygon

class ImageProcessor:
    def __init__(self, csvfile ,threshold=0.5):
        self.csvfile = csvfile
        data = pd.read_csv(csvfile)
        self.image = data.values #原始图像
        self.Ehimage = data.values #调整对比度后的
        self.width = self.image.shape[1]
        self.height = self.image.shape[0]
        self.threshold = threshold
        self.rectangle_points = [] 
        self.polygon_points = [] 
        #self.label_filename = 'labels.csv'
    
    def filter_strategy(self):
        # TODO
        pass
    
    def image_enhancement(self):
        # 图像增强的函数实现
        array_copy = np.copy(self.image)
        mean = np.mean(array_copy)
        array_copy[array_copy < mean] = 0
        max_value = np.max(array_copy)
        array_copy[array_copy > mean] = (array_copy[array_copy > mean] - mean) / (max_value - mean)
        kernel = np.ones((3, 3), np.uint8)
        for i in range(20):
            eroded_data = cv2.erode(array_copy, kernel)
        #plt.imshow(eroded_data, cmap='jet', vmin=0, vmax=1)
        # plt.show()
        #matplotlib.image.imsave(png_name[0]+'3.png', eroded_data, cmap='jet', vmin=0, vmax=1)
        self.Ehimage = eroded_data
    
    # 设置围挡区域外的值为固定值    
    def Set_fence_value(self):
        data = self.Ehimage 
        mask_inverse = 0.01
        # A cup
        if data.shape[0] == 288 and data.shape[1] == 769:
            x = [28, 133, 219, 290, 324, 397, 534, 635, 693, 750]
            y = [0, 119, 168, 192, 189, 189, 185, 140, 99, 0]
            mask = np.zeros_like(data, dtype=np.uint8)
            # 生成多边形掩码
            cv2.fillPoly(mask, [np.array(list(zip(x, y)), dtype=np.int32)], 255)
            # 将data转换为浮点型数组
            data = data.astype(float)
            # 将掩码外的像素设置为固定值mask_inverse
            # 假设mask_inverse是一个固定值
            data[np.where(mask == 0)] = mask_inverse
        # B cup
        if data.shape[0] == 384 and data.shape[1] == 897:
            x = [55, 120, 260, 437, 647, 741, 787, 792]
            y = [0, 151, 263, 310, 263, 193, 109, 0 ]
            mask = np.zeros_like(data, dtype=np.uint8)
            # 生成多边形掩码
            cv2.fillPoly(mask, [np.array(list(zip(x, y)), dtype=np.int32)], 255)
            # 将data转换为浮点型数组
            data = data.astype(float)
            # 将掩码外的像素设置为固定值mask_inverse
            data[np.where(mask == 0)] = mask_inverse     
        # C cup
        if data.shape[0] == 448 and data.shape[1] == 1089:
            x = [76, 149, 306, 502, 765, 881, 974, 1027]
            y = [0, 149, 308, 368, 325, 245, 142, 0 ]
            mask = np.zeros_like(data, dtype=np.uint8)
            # 生成多边形掩码
            cv2.fillPoly(mask, [np.array(list(zip(x, y)), dtype=np.int32)], 255)
            # 将data转换为浮点型数组
            data = data.astype(float)
            # 将掩码外的像素设置为固定值mask_inverse
            data[np.where(mask == 0)] = mask_inverse    
        # D cup
        if data.shape[0] == 544 and data.shape[1] == 1089:
            x = [59, 99, 208, 502, 743, 849, 945, 981]
            y = [0, 210, 359, 521, 471, 369, 250, 0 ]
            mask = np.zeros_like(data, dtype=np.uint8)
            # 生成多边形掩码
            cv2.fillPoly(mask, [np.array(list(zip(x, y)), dtype=np.int32)], 255)
            # 将data转换为浮点型数组
            data = data.astype(float)
            # 将掩码外的像素设置为固定值mask_inverse
            data[np.where(mask == 0)] = mask_inverse 
        # E cup        
        if data.shape[0] == 608 and data.shape[1] == 1089:
            x = [40, 1049, 1049,40]
            y = [1, 1, 470, 470 ]
            mask = np.zeros_like(data, dtype=np.uint8)
            # 生成多边形掩码
            cv2.fillPoly(mask, [np.array(list(zip(x, y)), dtype=np.int32)], 255)
            # 将data转换为浮点型数组
            data = data.astype(float)
            # 将掩码外的像素设置为固定值mask_inverse
            data[np.where(mask == 0)] = mask_inverse
        self.Ehimage = data
        # plt.imshow(self.Ehimage, cmap='jet', vmin=0, vmax=1)
        # plt.show()
        
    def show_image(self,image,points):
         # 绘制原始数据
        plt.imshow(self.image, cmap='jet',vmin=0, vmax=100)
        # 绘制边界框
        for i,region in enumerate(points):
            plt.plot(region[0], region[1], 'r-', linewidth=1)
            centroid_x = np.mean(region[0])
            centroid_y = np.mean(region[1])
            plt.text(centroid_x, centroid_y, str(i+1), color='r', fontsize=10, ha='center', va='center')
        # 显示图形
        plt.show()  

    def save_rectangle(self):
        # 保存矩形轮廓
        # 根据阈值创建二值图像
        binary_data = np.where(self.Ehimage >= self.threshold, 1, 0)
        # 标记连通区域
        labeled_data, num_labels = label(binary_data)
        # 找到每个连通区域的边界
        regions = find_objects(labeled_data)
        # 保存边界
        for region in regions:
            min_x, max_x = region[1].start, region[1].stop
            min_y, max_y = region[0].start, region[0].stop
            self.rectangle_points.append([[min_x, min_x, max_x, max_x, min_x],[min_y, max_y, max_y, min_y, min_y]]) 
        return self.rectangle_points
  
        

    def save_polygon(self):
        # 保存多边形轮廓
        scaled_data = np.uint8(self.Ehimage * 255)
        threshold_value = int(255*self.threshold)
        ret, thresh = cv2.threshold(scaled_data, threshold_value, 255, cv2.THRESH_BINARY)
        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_L1)
        # cv2.drawContours(scaled_data, contours, -1, (0, 255, 0), 2)
        # cv2.imshow('Thresholded Image', scaled_data)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()
        # 保存多边形
        # plt.imshow(self.Ehimage, cmap='jet',vmin=0, vmax=1)
        # plt.show() 
        filtered_contours = []
        for contour in contours:
            area = cv2.contourArea(contour)
            if len(contour) > 2 and area > 0:
                filtered_contours.append(contour)

        contours = filtered_contours
        if contours:
            for contour in contours:
            # 将每个轮廓转换为坐标列表形式
                contour = np.squeeze(contour)
                x = contour[:, 0]
                y = contour[:, 1]
                # 添加起始点，将轮廓闭合
                x = np.append(x, x[0])
                y = np.append(y, y[0])
                self.polygon_points.append([x,y])        
        return self.polygon_points



    
if __name__=='__main__':
    processer = ImageProcessor('202312040003_RCW_aD_2.csv',0.6)
    processer.image_enhancement()
    points = processer.save_rectangle()
    processer.show_image(processer.Ehimage,points)
    points = processer.save_polygon()
    processer.show_image(processer.Ehimage,points)



# file_name = '202312040003_LCW_aD_1.csv'
# png_name = file_name.split(".")

# # 读取CSV文件
# data = pd.read_csv(file_name)

# # 将DataFrame转换为二维数组
# array = data.values


    
# array_copy = np.copy(array)
# matplotlib.image.imsave(png_name[0]+'1.png', array_copy, cmap='jet', vmin=0, vmax=100)

# mean = np.mean(array_copy)
# array_copy[array_copy < mean] = 0
# max_value = np.max(array_copy)
# array_copy[array_copy > mean] = (array_copy[array_copy > mean] - mean) / (max_value - mean)

# matplotlib.image.imsave(png_name[0]+'2.png', array_copy, cmap='jet', vmin=0, vmax=1)

# kernel = np.ones((3, 3), np.uint8)
# # 对归一化后的数据进行腐蚀操作
# for i in range(20):
#     eroded_data = cv2.erode(array_copy, kernel)
# matplotlib.image.imsave(png_name[0]+'3.png', eroded_data, cmap='jet', vmin=0, vmax=1)

# # histogram = plt.hist(eroded_data.flatten(), bins=100)
# # # 显示直方图
# # plt.show()

# ## 显示矩形框
# # threshold = 0.5
# # # 根据阈值创建二值图像
# # binary_data = np.where(eroded_data >= threshold, 1, 0)
# # # 标记连通区域
# # labeled_data, num_labels = label(binary_data)
# # # 找到每个连通区域的边界
# # regions = find_objects(labeled_data)
# # # 绘制原始数据
# # plt.imshow(eroded_data, cmap='jet')

# # # 绘制边界框
# # for region in regions:
# #     min_x, max_x = region[1].start, region[1].stop
# #     min_y, max_y = region[0].start, region[0].stop
# #     plt.plot([min_x, min_x, max_x, max_x, min_x], [min_y, max_y, max_y, min_y, min_y], 'r-', linewidth=2)
    
# # # 显示图形
# # plt.show()


# # 将数据进行缩放到0-255的范围
# scaled_data = np.uint8(eroded_data * 255)

# # 2. 图像分割
# # 转换为图像格式
# image = scaled_data

# # 自定义阈值
# threshold_value = int(255*0.7)

# # 阈值分割
# ret, thresh = cv2.threshold(image, threshold_value, 255, cv2.THRESH_BINARY)

# # 3. 轮廓检测
# contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_L1)

# # # 4. 边界拟合
# # for contour in contours:
# #     # 计算边界拟合
# #     epsilon = 0.03 * cv2.arcLength(contour, True)
# #     approx = cv2.approxPolyDP(contour, epsilon, True)

# #     # 绘制边界
# #     cv2.drawContours(image, [approx], 0, (0, 255, 0), 2)
# # 绘制轮廓
# cv2.drawContours(image, contours, -1, (0, 255, 0), 2)


# # 5. 可视化
# cv2.imshow("Result", scaled_data)
# cv2.waitKey(0)
# cv2.destroyAllWindows()

